from model import *
from utils import *

batch_size = 100
num_epochs = 30
learning_rate = 1e-4

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

train_iter, test_iter = load_data_mnist(batch_size)
model = CVAE()
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    train_loss = 0
    i = 0
    for batch_id, data in enumerate(train_iter):
        img, label = data
        img, label = img.to(device), label.to(device)
        inputs = img.reshape(img.shape[0], -1)
        label = torch.zeros(len(label), 10).scatter(1, label.unsqueeze(1), 1)

        recon, mean, std = model(inputs, label)
        loss = model.vae_loss(inputs, recon, mean, std)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        i += 1

    if (epoch + 1) % 5 == 0:
        print("Epoch[{}/{}], Train_loss:{:.6f}".format(
            epoch+1, num_epochs, train_loss/i))

torch.save(model.state_dict(), 'cvae.params')
